#include "BufferManager.h"
#include <iostream>
#include <fstream>
#include <cstring>

BufferManager::BufferManager(VulkanApplication* app)
    : _descriptorPool(nullptr), _descriptorSetLayout(nullptr), _application(app)
{}

BufferManager::~BufferManager()
{}

bool BufferManager::prepare(const std::string& vertFile, const std::string& fragFile)
{
    VkVertexInputBindingDescription bindingDescription
    {
        0, sizeof(Vertex),
        VK_VERTEX_INPUT_RATE_VERTEX
    };

    std::vector<VkVertexInputAttributeDescription> attrDescription;
    attrDescription.push_back(
        {
            0, 0,  // location, binding
            VK_FORMAT_R32G32B32_SFLOAT,  // format
            0                         // offset
        });
    attrDescription.push_back(
        {
            1, 0,  // location, binding
            VK_FORMAT_R32G32B32_SFLOAT,  // format
            sizeof(float) * 3         // offset
        });

    VkPipelineVertexInputStateCreateInfo inputVertexInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        nullptr, 0,
        1, &bindingDescription,
        (uint32_t)attrDescription.size(), attrDescription.data()
    };

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
        nullptr, 0,
        VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, VK_FALSE
    };
    
    // Create graphics pipeline
    return _application->setUpPipeline(
        vertFile, fragFile, inputVertexInfo, inputAssemblyInfo, _descriptorSetLayout);
}

bool BufferManager::createDescriptorSet()
{
    // Create uniform desciption layout
    VkDescriptorSetLayoutBinding uboLayoutBinding
    {
        0,  // binding
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        1,  // descriptorCount
        VK_SHADER_STAGE_VERTEX_BIT,
        nullptr
    };

    VkDescriptorSetLayoutCreateInfo layoutInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        nullptr, 0,
        1, &uboLayoutBinding
    };

    VkResult result = vkCreateDescriptorSetLayout(
        _application->getDevice(), &layoutInfo, nullptr, &_descriptorSetLayout);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description-set layout.\n");
        return false;
    }

    // Create uniform desciption pool
    VkDescriptorPoolSize poolSize
    { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 };

    VkDescriptorPoolCreateInfo poolInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        nullptr, 0,
        1, 1, &poolSize
    };

    result = vkCreateDescriptorPool(
        _application->getDevice(), &poolInfo, nullptr, &_descriptorPool);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description pool.\n");
        return false;
    }

    // Create uniform description sets
    VkDescriptorSetAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        nullptr,
        _descriptorPool,
        1, &_descriptorSetLayout
    };

    result = vkAllocateDescriptorSets(_application->getDevice(), &allocInfo, &_descriptorSet);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description set.\n");
        return false;
    }
    return true;
}

void BufferManager::cleanup()
{
    for (std::map<std::string, VkBuffer>::iterator itr = _bufferMap.begin();
        itr != _bufferMap.end(); ++itr)
    {
        vkDestroyBuffer(_application->getDevice(), itr->second, nullptr);
    }

    for (std::map<std::string, VkDeviceMemory>::iterator itr = _bufferMemoryMap.begin();
         itr != _bufferMemoryMap.end(); ++itr)
    {
        vkFreeMemory(_application->getDevice(), itr->second, nullptr);
    }

    vkDestroyDescriptorSetLayout(_application->getDevice(), _descriptorSetLayout, nullptr);
    vkDestroyDescriptorPool(_application->getDevice(), _descriptorPool, nullptr);
}

VkBuffer BufferManager::createVertexBuffer(const std::string& name, const std::vector<Vertex>& vertices)
{
    size_t dataSize = sizeof(Vertex) * vertices.size();
    VkBuffer vertexBuffer = nullptr;
    VkDeviceMemory vertexBufferMemory = nullptr;
    if (!createBuffer(vertexBuffer, dataSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), vertexBuffer, &memRequirements);
    if (!createBufferMemory(vertexBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    vkBindBufferMemory(_application->getDevice(), vertexBuffer, vertexBufferMemory, 0);
    _bufferMap[name] = vertexBuffer;
    _bufferMemoryMap[name] = vertexBufferMemory;

    void* data = nullptr;
    vkMapMemory(_application->getDevice(), vertexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, vertices.data(), dataSize);
    vkUnmapMemory(_application->getDevice(), vertexBufferMemory);
    return vertexBuffer;
}

VkBuffer BufferManager::createIndexBuffer(const std::string& name, const std::vector<unsigned int>& indices)
{
    size_t dataSize = sizeof(unsigned int) * indices.size();
    VkBuffer indexBuffer = nullptr;
    VkDeviceMemory indexBufferMemory = nullptr;
    if (!createBuffer(indexBuffer, dataSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), indexBuffer, &memRequirements);
    if (!createBufferMemory(indexBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    vkBindBufferMemory(_application->getDevice(), indexBuffer, indexBufferMemory, 0);
    _bufferMap[name] = indexBuffer;
    _bufferMemoryMap[name] = indexBufferMemory;

    void* data = nullptr;
    vkMapMemory(_application->getDevice(), indexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, indices.data(), dataSize);
    vkUnmapMemory(_application->getDevice(), indexBufferMemory);
    return indexBuffer;
}

VkBuffer BufferManager::createUniformBuffer(const std::string& name, size_t uboSize)
{
    VkBuffer uniformBuffer = nullptr;
    VkDeviceMemory uniformBufferMemory = nullptr;
    if (!createBuffer(uniformBuffer, uboSize, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), uniformBuffer, &memRequirements);
    if (!createBufferMemory(uniformBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;
    
    vkBindBufferMemory(_application->getDevice(), uniformBuffer, uniformBufferMemory, 0);
    _bufferMap[name] = uniformBuffer;
    _bufferMemoryMap[name] = uniformBufferMemory;

    // Update description set
    VkDescriptorBufferInfo bufferInfo
    { uniformBuffer, 0, uboSize };

    VkWriteDescriptorSet descriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        _descriptorSet, 0,  // dstSet, dstBinding
        0, 1,              // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        nullptr, &bufferInfo, nullptr
    };
    vkUpdateDescriptorSets(_application->getDevice(), 1, &descriptorWrite, 0, nullptr);
    return uniformBuffer;
}

void* BufferManager::mapBufferMemory(const std::string& name, size_t dataSize)
{
    void* dataPtr = nullptr;
    if (_bufferMemoryMap.find(name) != _bufferMemoryMap.end())
    {
        vkMapMemory(_application->getDevice(), _bufferMemoryMap[name],
                    0, dataSize, 0, &dataPtr);
    }
    return dataPtr;
}

void BufferManager::unmapBufferMemory(const std::string& name)
{
    if (_bufferMemoryMap.find(name) != _bufferMemoryMap.end())
        vkUnmapMemory(_application->getDevice(), _bufferMemoryMap[name]);
}

VkBuffer BufferManager::getBuffer(const std::string& name)
{
    std::map<std::string, VkBuffer>::iterator itr = _bufferMap.find(name);
    if (itr != _bufferMap.end()) return itr->second;
    else return nullptr;
}

bool BufferManager::createBuffer(VkBuffer& buffer, size_t dataSize, VkBufferUsageFlags usage)
{
    VkBufferCreateInfo bufferInfo
    {
        VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        nullptr, 0,
        dataSize,
        usage,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr  // pQueueFamilyIndices
    };

    VkResult result = vkCreateBuffer(_application->getDevice(), &bufferInfo, nullptr, &buffer);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create data buffer.\n");
        return false;
    }
    return true;
}

bool BufferManager::createBufferMemory(VkDeviceMemory& bufferMemory,
                                       VkMemoryRequirements& memRequirements, uint32_t propertyBits)
{
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(_application->getPhysicalDevice(), &memProperties);

    uint32_t memoryTypeIndex = 0;
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; ++i)
    {
        if ((memRequirements.memoryTypeBits & (1 << i)) &&
            (memProperties.memoryTypes[i].propertyFlags & propertyBits) == propertyBits)
        { memoryTypeIndex = i; }
    }

    VkMemoryAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        nullptr,
        memRequirements.size,
        memoryTypeIndex
    };

    VkResult result = vkAllocateMemory(_application->getDevice(), &allocInfo, nullptr, &bufferMemory);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create vertex buffer memory.\n");
        return false;
    }
    return true;
}
